import os
import time
import torch
import numpy as np
import utils
import logging
from collections import defaultdict

from options import *
from model.hidden import Hidden
from average_meter import AverageMeter
from options import HiDDenConfiguration, TrainingOptions

def train(model: Hidden,
          device: torch.device,
          hidden_config: HiDDenConfiguration,
          train_options: TrainingOptions,
          this_run_folder: str,
          tb_logger):
    """
    Trains the HiDDeN model
    :param model: The model
    :param device: torch.device object, usually this is GPU (if avaliable), otherwise CPU.
    :param hidden_config: The network configuration
    :param train_options: The training settings
    :param this_run_folder: The parent folder for the current training run to store training artifacts/results/logs.
    :param tb_logger: TensorBoardLogger object which is a thin wrapper for TensorboardX logger.
                Pass None to disable TensorboardX logging
    :return:
    """
    # torch.autograd.set_detect_anomaly(True)
    train_data, val_data = utils.get_data_loaders(hidden_config, train_options)
    file_count = len(train_data.dataset)
    if file_count % train_options.batch_size == 0:
        steps_in_epoch = file_count // train_options.batch_size
    else:
        steps_in_epoch = file_count // train_options.batch_size + 1

    attack = True
    print_each = 10
    images_to_save = 8
    saved_images_size = (512, 512)

    # 加载预训练的噪声层参数
    if attack:
        # 加载模型时，使用map_location参数将模型映射到当前设备 在设备1上训练的模型，不能直接加载到设备0中，总之不能在不同设备上运行
        attack_state_dict = torch.load('/home/zzc/Experience/HiDDeN_zzc/attack_runs/run-Brightness_2023-12-28_23-17-16/model/best_model.pth')
        # attack_state_dict = torch.load('/home/zzc/deeplearn_test/HiDDeN/attack_runs/run-Gaussian(0.06)_2023-12-28_22-52-49/model/best_model.pth', map_location=torch.device(device))
        # attack_state_dict = torch.load('/home/zzc/deeplearn_test/HiDDeN/attack_runs/run-salt_2023-12-28_22-54-10/model/best_model.pth', map_location=torch.device(device))
        # attack_state_dict = torch.load('/home/zzc/deeplearn_test/HiDDeN/attack_runs/run-Contrast_2023-12-28_23-20-33/model/best_model.pth', map_location=torch.device(device))
        # attack_state_dict = torch.load('/home/zzc/deeplearn_test/HiDDeN/attack_runs/run-Hue_2023-12-28_23-22-56/model/best_model.pth', map_location=torch.device(device))
        # attack_state_dict = torch.load('/home/zzc/Experience/HiDDeN_zzc/attack_runs/run-JPEG(50)_2023-12-28_22-48-48/model/best_model.pth', map_location=torch.device(device))
        # attack_state_dict = torch.load('/home/zzc/deeplearn_test/HiDDeN/attack_runs/run-Saturation_2023-12-28_23-21-25/model/best_model.pth', map_location=torch.device(device))
        # attack_state_dict = torch.load('/home/zzc/deeplearn_test/HiDDeN/attack_runs/run-Quantization_2023-12-28_23-00-07/model/best_model.pth', map_location=torch.device(device))
        # attack_state_dict = torch.load('/home/zzc/Experience/HiDDeN/attack_runs/run-JPEG(50)_Mask_2024-02-27_17-34-52/model/best_model.pth', map_location=torch.device(device))
        model.encoder_decoder.noiser.load_state_dict(attack_state_dict)
    # 固定噪声层的参数，不进行梯度更新
    for param in model.encoder_decoder.noiser.parameters():
        param.requires_grad = False

    for epoch in range(train_options.start_epoch, train_options.number_of_epochs + 1):
        logging.info('\nStarting epoch {}/{}'.format(epoch, train_options.number_of_epochs))
        logging.info('Batch size = {}\nSteps in epoch = {}'.format(train_options.batch_size, steps_in_epoch))
        training_losses = defaultdict(AverageMeter)
        epoch_start = time.time()
        step = 1


        for image, _ in train_data:
            image = image.to(device)
            message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device)
            losses, _ = model.train_on_batch([image, message])

            for name, loss in losses.items():
                training_losses[name].update(loss)
            if step % print_each == 0 or step == steps_in_epoch:
                logging.info(
                    'Epoch: {}/{} Step: {}/{}'.format(epoch, train_options.number_of_epochs, step, steps_in_epoch))
                utils.log_progress(training_losses)
                logging.info('-' * 40)
            step += 1

        train_duration = time.time() - epoch_start
        logging.info('Epoch {} training duration {:.2f} sec'.format(epoch, train_duration))
        logging.info('-' * 40)
        utils.write_losses(os.path.join(this_run_folder, 'train.csv'), training_losses, epoch, train_duration)
        if tb_logger is not None:
            tb_logger.save_losses(training_losses, epoch)
            tb_logger.save_grads(epoch)
            tb_logger.save_tensors(epoch)

        first_iteration = True
        validation_losses = defaultdict(AverageMeter)
        logging.info('Running validation for epoch {}/{}'.format(epoch, train_options.number_of_epochs))
        for image, _ in val_data:
            image = image.to(device)
            message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device)
            losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch([image, message])
            for name, loss in losses.items():
                validation_losses[name].update(loss)
            if first_iteration:
                if hidden_config.enable_fp16:
                    image = image.float()
                    encoded_images = encoded_images.float()
                revise_image = torch.abs(image - encoded_images) * 5
                revise_image_ = torch.abs(image - noised_images) * 5
                utils.save_images(image.cpu()[:images_to_save, :, :, :],
                                        encoded_images[:images_to_save, :, :, :].cpu(),
                                        noised_images[:images_to_save, :, :, :].cpu(),
                                        revise_image[:images_to_save, :, :, :].cpu(),
                                        revise_image_[:images_to_save, :, :, :].cpu(),
                                        epoch,
                                        os.path.join(this_run_folder, 'images'), resize_to=saved_images_size)
                first_iteration = False

        utils.log_progress(validation_losses)
        logging.info('-' * 40)
        utils.save_checkpoint(model, train_options.experiment_name, epoch, os.path.join(this_run_folder, 'checkpoints'))
        utils.write_losses(os.path.join(this_run_folder, 'validation.csv'), validation_losses, epoch,
                           time.time() - epoch_start)
        # 将参数信息保存到字符串中
        param_info = ""
        for name, param in model.encoder_decoder.named_parameters():
            param_info += f"Parameter name: {name}, Shape: {param.shape}\n"
            param_info += f"Value: {param}\n"

        # 将参数信息写入到txt文件中
        with open('model_params.txt', 'w') as f:
            f.write(param_info)

        print("Model parameters info saved to 'model_params.txt'")
